/*
 * Description: 资源泄露相关事件上报
 */
#include "resource_event_handler.h"
#include <windows.h>
#include <psapi.h>
#include <iostream>
#include <list>
#include <map>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <chrono>
#include <vector>
#include <sstream>
#include <winver.h>
#include <codecvt>
#include "event_reporter.h"
#include "stability_common.h"
#include "singleton.h"
#include "log.h"
#include "text_util/string_proc.h"
#pragma comment(lib, "version.lib")

namespace DFX {
wostream &operator << (wostream &os, const ResourceInfo &ri)
{
    os << L"Process ID: " << ri.processId << L";" << L"Process Name: " << ri.proName << L";" << L"Process Version: " <<
        ri.proVersion << L";" << L"Resource Type: " << static_cast<int>(ri.resouceType) << L";" << L"Has Report: " <<
        (ri.hasReport ? L"Yes" : L"No") << L";" << L"Last Error Flag: " << (ri.lastErrorFlag ? L"Yes" : L"No") <<
        L";" << L"Current Resource Value: " << ri.resouceValule << L";" << L"Resource Value History:";

    for (auto val : ri.resouceValuleHistory) {
        os << L"  " << val << L",";
    }
    os << L"\n";
    return os;
}

void ResourceEventHandler::CreateSysInforBuffer(ULONG bufferSize)
{
    if ((bufferSize == 0) || (bufferSize > MAX_SYS_INFORMATION_IN)) {
        return;
    }
    if (sysInfoBuffer != nullptr) {
        delete[] sysInfoBuffer;
        sysInfoBuffer = nullptr;
    }
    sysInfoBuffer = new (nothrow) char[bufferSize];
    if (sysInfoBuffer == nullptr) {
        LOG_ERR("[DFX - Stability] - new memory for NtQuerySystemInformation buffer failed,err={}", GetLastError());
        return;
    }
    sysInfoBufferLen = bufferSize;
}

bool ResourceEventHandler::QuerySystemProcessInfomation()
{
    ULONG needLen = 0;
    SecureZeroMemory(sysInfoBuffer, sysInfoBufferLen);
    LONG status = NtQuerySystemInformation(SYSTEM_PROCESS_INFORMATION, sysInfoBuffer, sysInfoBufferLen, &needLen);
    if (status == STATUS_INFO_LENGTH_MISMATCH) {
        LOG_ERR("[DFX - Stability] - buff not enough sysInfoBufferLen={}, needLen={}", sysInfoBufferLen, needLen);
        CreateSysInforBuffer(needLen);
        if (sysInfoBuffer == nullptr) {
            return false;
        }
        status = NtQuerySystemInformation(SYSTEM_PROCESS_INFORMATION, sysInfoBuffer, sysInfoBufferLen, &needLen);
    }
    if (status != 0) {
        LOG_ERR("[DFX - Stability] - NtQuerySystemInformation failed, status={}", status);
        return false;
    }
    return true;
}

void ResourceEventHandler::GetResource(list<ResourceInfo> &faultInfo)
{
    if (!QuerySystemProcessInfomation()) {
        return;
    }
    ULONG processOffset = 0;
    map<wstring, int64_t> tempProcessCount;
    map<wstring, HANDLE> tempProcessId;
    while (true) {
        char *processInfoBuff = reinterpret_cast<char *>(sysInfoBuffer) + processOffset;
        SystemProcess *processInfo = reinterpret_cast<SystemProcess *>(processInfoBuff);
        wstring proname =
            wstring(processInfo->processName.buffer, processInfo->processName.length / sizeof(WCHAR));

        if (processInfo->handleCount > handleLeakThreshold && handleLeakThreshold != 0) { // handle leak
            faultInfo.push_back(
                ResourceInfo(processInfo->processId, proname, RESOUCE_TYPE_HANDLE, processInfo->handleCount));
        }

        if ((processInfo->workingSetSize / BUFFER_LEN_1M) > memoryLeakThreshold && memoryLeakThreshold != 0) {
            faultInfo.push_back(ResourceInfo(processInfo->processId, proname, RESOUCE_TYPE_MEMORY,
                static_cast<int64_t>(processInfo->workingSetSize)));
        }

        if (processInfo->numberOfThread >= 1) {
            if (tempProcessCount.find(proname) != tempProcessCount.end()) { // process count check
                ((tempProcessCount.find(proname))->second)++;
            } else {
                tempProcessCount.insert(make_pair(proname, 1));
                tempProcessId.insert(make_pair(proname, processInfo->processId));
            }
        }

        if (processInfo->numberOfThread > threadLeakThreshold && threadLeakThreshold != 0) { // thread leak
            faultInfo.push_back(
                ResourceInfo(processInfo->processId, proname, RESOUCE_TYPE_THREAD, processInfo->numberOfThread));
        }
        processOffset = processOffset + processInfo->nextEntryOffset;
        if (processInfo->nextEntryOffset == 0) {
            break;
        }
    }

    for (auto item : tempProcessCount) {
        if ((UINT64)(item.second) > processLeakThreshold && processLeakThreshold != 0) {
            auto processIdItem = tempProcessId.find(item.first);
            HANDLE processId = processIdItem != tempProcessId.end() ? processIdItem->second : 0;
            faultInfo.push_back(ResourceInfo(processId, item.first, RESOUCE_TYPE_PROCESS, item.second));
        }
    }
}

wstring ResourceEventHandler::GetFileVersion(const wstring &filePath)
{
    DWORD dummyHandle;
    DWORD versionInfoSize = GetFileVersionInfoSizeW(filePath.c_str(), &dummyHandle);
    if (versionInfoSize == 0) {
        LOG_ERR("[DFX - Stability] - Error getting version info size");
        return L"";
    }

    LPVOID versionInfoBuffer = new (std::nothrow) BYTE[versionInfoSize];
    if (versionInfoBuffer == nullptr) {
        LOG_ERR("[DFX - Stability] - VersionInfoBuffer is NULL");
        return L"";
    }
    if (!GetFileVersionInfoW(filePath.c_str(), 0, versionInfoSize, versionInfoBuffer)) {
        delete[] versionInfoBuffer;
        LOG_ERR("[DFX - Stability] - Error getting version info size");
        return L"";
    }

    UINT length;
    VS_FIXEDFILEINFO *fixedFileInfo;
    if (!VerQueryValueW(versionInfoBuffer, L"\\", (LPVOID *)&fixedFileInfo, &length)) {
        delete[] versionInfoBuffer;
        LOG_ERR("[DFX - Stability] - Error querying version value");
        return L"";
    }

    wstringstream wss;
    wss << ((fixedFileInfo->dwFileVersionMS >> VERSION_MAJOR_SHIFT) & 0xFFFF) << L"." <<
        ((fixedFileInfo->dwFileVersionMS) & 0xFFFF) <<
        L"." << ((fixedFileInfo->dwFileVersionLS >> VERSION_MAJOR_SHIFT) & 0xFFFF) << L"." <<
        ((fixedFileInfo->dwFileVersionLS) & 0xFFFF);

    wstring versionStr = wss.str();
    delete[] versionInfoBuffer;
    return versionStr;
}

wstring ResourceEventHandler::GetProcVersionInfo(HANDLE processId)
{
    if (processId == nullptr) {
        return L"";
    }
    auto processIdAddress = reinterpret_cast<uintptr_t>(processId);
    HANDLE process = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, static_cast<DWORD>(processIdAddress));
    if (process == nullptr) {
        return L"";
    }

    wstring imagePath = L"";
    WCHAR fileName[MAX_PATH] = { 0 };
    DWORD fileNameLen = GetModuleFileNameExW(process, nullptr, fileName, MAX_PATH);
    if (fileNameLen > 0) {
        imagePath = wstring(fileName, fileNameLen);
    } else {
        fileNameLen = MAX_PATH;
        if (QueryFullProcessImageNameW(process, 0, fileName, &fileNameLen)) {
            imagePath = wstring(fileName, fileNameLen);
        }
    }
    CloseHandle(process);
    process = nullptr;
    string fileVersion = "";
    return GetFileVersion(imagePath);
}

int ResourceEventHandler::RemoveRepairResouce()
{
    for (auto item = handleErrorResource.begin(); item != handleErrorResource.end();) {
        if (!item->second.lastErrorFlag) {
            handleErrorResource.erase(item++);
            continue;
        }
        item->second.lastErrorFlag = false;
        ++item;
    }
    for (auto item = memErrorResource.begin(); item != memErrorResource.end();) {
        if (!item->second.lastErrorFlag) {
            memErrorResource.erase(item++);
            continue;
        }
        item->second.lastErrorFlag = false;
        ++item;
    }
    for (auto item = threadErrorResource.begin(); item != threadErrorResource.end();) {
        if (!item->second.lastErrorFlag) {
            threadErrorResource.erase(item++);
            continue;
        }
        item->second.lastErrorFlag = false;
        ++item;
    }
    for (auto process = processFaultResource.begin(); process != processFaultResource.end();) {
        if (!process->second.lastErrorFlag) {
            processFaultResource.erase(process++);
            continue;
        }
        process->second.lastErrorFlag = false;
        ++process;
    }
    return 0;
}


void ResourceEventHandler::GetResouceLeakData(vector<map<string, string>> &eventDatas, ResourceInfo &info)
{
    map<string, string> eventData;
    if (static_cast<size_t>(info.resouceType) > RESOURCE_EVENT_NAME.size()) {
        LOG_ERR("[DFX - Stability] - resouceType error, resouceType={}", info.resouceType);
        return;
    }
    eventData["EventName"] = RESOURCE_EVENT_NAME[info.resouceType];
    eventData["ProcessName"] = StringProc::WStringToString(info.proName);
    eventData["ProcessVersion"] = StringProc::WStringToString(info.proVersion);
    eventData["LeakInformation"] = to_string(info.resouceValule);
    if (!eventData["ProcessName"].empty()) {
        eventDatas.push_back(eventData);
    } else {
        LOG_WARN("[DFX - Stability] - ProcessName is empty");
    }
    info.hasReport = true;
}

vector<map<string, string>> ResourceEventHandler::GetAllLeakEvents()
{
    vector<map<string, string>> eventDatas;
    map<HANDLE, ResourceInfo>::iterator item;
    for (item = handleErrorResource.begin(); item != handleErrorResource.end(); ++item) {
        if (!item->second.hasReport && item->second.resouceValuleHistory.size() > leakTimesThreshold) {
            GetResouceLeakData(eventDatas, item->second);
        }
    }
    for (item = memErrorResource.begin(); item != memErrorResource.end(); ++item) {
        if (!item->second.hasReport && item->second.resouceValuleHistory.size() > leakTimesThreshold) {
            GetResouceLeakData(eventDatas, item->second);
        }
    }
    for (item = threadErrorResource.begin(); item != threadErrorResource.end(); ++item) {
        if (!item->second.hasReport && item->second.resouceValuleHistory.size() > leakTimesThreshold) {
            GetResouceLeakData(eventDatas, item->second);
        }
    }

    map<wstring, ResourceInfo>::iterator process;
    for (process = processFaultResource.begin(); process != processFaultResource.end(); ++process) {
        if (!process->second.hasReport && process->second.resouceValuleHistory.size() > leakTimesThreshold) {
            GetResouceLeakData(eventDatas, item->second);
        }
    }
    return eventDatas;
}

void ResourceEventHandler::CheckResouceImpl(ResourceInfo srcInfo, map<HANDLE, ResourceInfo> &targetInfo)
{
    map<HANDLE, ResourceInfo>::iterator history = targetInfo.find(srcInfo.processId);
    if (history != targetInfo.end()) {
        history->second.lastErrorFlag = true;
        history->second.resouceValuleHistory.push_back(srcInfo.resouceValule);
    } else {
        srcInfo.lastErrorFlag = true;
        targetInfo.insert(make_pair(srcInfo.processId, srcInfo));
    }
}

void ResourceEventHandler::CheckResouce(list<ResourceInfo> &faultInfo)
{
    for (auto item : faultInfo) {
        item.proVersion = GetProcVersionInfo(item.processId);
        switch (item.resouceType) {
            case RESOUCE_TYPE_HANDLE: {
                CheckResouceImpl(item, handleErrorResource);
                break;
            }
            case RESOUCE_TYPE_PROCESS: {
                auto history = processFaultResource.find(item.proName);
                if (history != processFaultResource.end()) {
                    history->second.lastErrorFlag = true;
                    history->second.resouceValuleHistory.push_back(item.resouceValule);
                } else {
                    item.lastErrorFlag = true;
                    processFaultResource.insert(make_pair(item.proName, item));
                }
                break;
            }
            case RESOUCE_TYPE_MEMORY: {
                CheckResouceImpl(item, memErrorResource);
                break;
            }
            case RESOUCE_TYPE_THREAD: {
                CheckResouceImpl(item, threadErrorResource);
                break;
            }
            default:
                break;
        }
    }
    RemoveRepairResouce();
}

bool ResourceEventHandler::Init()
{
    runEvent = CreateEvent(nullptr, TRUE, FALSE, nullptr);
    if (runEvent == nullptr) {
        LOG_ERR("[DFX - Stability] - create event for thread control failed,err={}", GetLastError());
        return false;
    }
    CreateSysInforBuffer(sysInfoBufferLen);
    if (sysInfoBuffer == nullptr) {
        LOG_ERR("[DFX - Stability] - VirtualAlloc for NtQuerySystemInformation buffer failed,err={}", GetLastError());
        return false;
    }
    ntDllModule = GetModuleHandleW(NTDLL_NAME.c_str());
    if (ntDllModule == nullptr) {
        LOG_ERR("[DFX - Stability] - GetModuleHandle for ntdll.dll failed,err={}", GetLastError());
        return false;
    }
    NtQuerySystemInformation = (FunNtQuerySystemInformation)GetProcAddress(ntDllModule, QUERY_SYSTEM_FUN_NAME.c_str());
    if (NtQuerySystemInformation == nullptr) {
        LOG_ERR("[DFX - Stability] - GetProcAddress for NtQuerySystemInformation failed,err={}", GetLastError());
        return false;
    }
    return true;
}

void ResourceEventHandler::DeInit()
{
    if (sysInfoBuffer != nullptr) {
        delete[] sysInfoBuffer;
        sysInfoBuffer = nullptr;
    }
    if (ntDllModule != nullptr) {
        FreeLibrary(ntDllModule);
        ntDllModule = nullptr;
    }
    if (runEvent != nullptr) {
        CloseHandle(runEvent);
        runEvent = nullptr;
    }
}

ResourceEventHandler::~ResourceEventHandler()
{
    DeInit();
}

void ResourceEventHandler::ReportResEvents()
{
    list<ResourceInfo> faultInfo;
    GetResource(faultInfo);
    LOG_INFO("[DFX - Stability] - faultInfo size: {}", faultInfo.size());
    CheckResouce(faultInfo);
    auto events = GetAllLeakEvents();
    LOG_INFO("[DFX - Stability] - AllLeakEvents size: {}", events.size());
    for (auto event : events) {
        auto iter = event.find("EventName");
        if (iter != event.end()) {
            string eventName = iter->second;
            event.erase(iter);
            EVENT_REPORTER.ReportEventData(eventName, event);
        } else {
            LOG_ERR("[DFX - Stability] - ReportResEvents failed, eventName is null");
        }
    }
}
}